#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <assert.h>

#include "graph.h"
#include "statistics.h"
#include "sample.h"

static void mcmc_rand_init(Graph_t * p_graph, int num_insts,
                           Instance_t * p_insts)
{
  int i;
  
  for (i=0; i<num_insts; i++)
  {
    Instance_t * p_inst;
    int v;
    
    p_inst = p_insts + i;
    
    for (v=0; v<p_graph->graph_num_vars; v++)
    {
      Variable_t * p_var;

      if (p_inst->inst_p_is_obs[v])
        continue;
      
      p_var = p_graph->graph_p_vars + v;
      
      p_inst->inst_p_val[v] = rand() % p_var->var_card;
    }
  }
}

/* go through each variable and sample a value for it using Gibbs */
static void mcmc_kernel_1var(Graph_t * p_graph, Instance_t * p_inst,
                             double * p_weights)
{
  int vi;                                    /* variable index */

  /* scan each non-observed variable */
  /* IMPROVEME: keep a list of non-observed variables */
  for (vi=0; vi<p_graph->graph_num_vars; vi++)
  {
    Variable_t * p_var;
    int value;
        
    if (p_inst->inst_p_is_obs[vi])
      continue;
        
    p_var = p_graph->graph_p_vars + vi;
        
    /* compute a weight for each possible value of the variable */
    for (value=0; value<p_var->var_card; value ++)
    {
      int fi;                            /* factor index */
          
      p_inst->inst_p_val[vi] = value;
      /* compute the weight for this value */
      p_weights[value] = 0;

      for (fi=0; fi<p_var->var_num_facs; fi++)
      {
        Factor_t * p_fac;
            
        p_fac = p_var->var_p_facs[fi];

        p_weights[value] += graph_factor(p_fac, p_inst->inst_p_val);
      }
    }

    /* randomly choose a value proportional to the weights */
    value = sample_log_weights(p_var->var_card, p_weights);

    p_inst->inst_p_val[vi] = value;
  }
}

#ifdef DEBUG
static int debug = 1;
#endif

#ifdef KERNEL_NBR_VAR
static void mcmc_kernel_nbr_var(Graph_t * p_graph, Instance_t * p_inst,
                                double * p_weights)
{
  int vi;                                    /* variable index */

  /* scan each non-observed variable */
  /* IMPROVEME: keep a list of non-observed variables */
  for (vi=0; vi<p_graph->graph_num_vars; vi++)
  {
    Variable_t * p_var;
    int nbri;
    int nbr_card;
    int index;
    int is_evi_mod;  
        
    if (p_inst->inst_p_is_obs[vi])
      continue;
    
    p_var = p_graph->graph_p_vars + vi;

    /* compute the cardinality of all the neigbors together */
    nbr_card = 1;
    for (nbri = 0; nbri < p_var->var_num_nbr_vars; nbri ++)
    {
      Variable_t * p_nbr_var;
      
      p_nbr_var = p_var->var_p_nbr_vars[nbri];
      
      nbr_card *= p_nbr_var->var_card;
      
      if (nbr_card > MAX_CARD)
        break;
    }

    if (nbr_card > MAX_CARD)
      continue;

#ifdef DEBUG
    if (debug)
    {
      int v;
      printf("inst:");
      for (v=0; v<p_graph->graph_num_vars; v++)
        printf(" %d", p_inst->inst_p_val[v]);
      printf("\n");
    }
#endif

    /* compute the weight of each possible index */
    /* IMPROVEME: this is a pathetically naive n^2 algorithm, it can be
     * improved by making a pass through all the variables incrementing their
     * value and only updating the weight of factors related to the modified
     * variable */
    for (index = 0; index < nbr_card; index ++)
    {
      is_evi_mod = graph_set_vars(p_var->var_num_nbr_vars,
                                  p_var->var_p_nbr_vars, index, p_inst);
      
      /* if the evidence is modified by this index then skip it */
      if (is_evi_mod)
      {
        p_weights[index] = -HUGE_VAL;
        continue;
      }
      
      p_weights[index] = 0;
      for (nbri = 0; nbri < p_var->var_num_nbr_facs; nbri ++)
      {
        Factor_t * p_nbr_fac;
        
        p_nbr_fac = p_var->var_p_nbr_facs[nbri];

        p_weights[index] += graph_factor(p_nbr_fac, p_inst->inst_p_val);
      }
    }
    
    index = sample_log_weights(nbr_card, p_weights);

#ifdef DEBUG    
    if (debug)
    {
      int i;
      printf("Sampling %d's nbrs (max card %d):", vi, nbr_card);
      for (i = 0; i < nbr_card; i ++)
        printf(" %g", p_weights[i]);
      printf(" --> %d", index);
      printf("\n");
    }
#endif

    is_evi_mod = graph_set_vars(p_var->var_num_nbr_vars, p_var->var_p_nbr_vars,
                                index, p_inst);

    /* the evidence should not be modified since any index which does
     * so should have a zero weight */
    assert(!is_evi_mod);
  }

#ifdef DEBUG
  if (debug)
  {
    int v;
    printf("inst:");
    for (v=0; v<p_graph->graph_num_vars; v++)
      printf(" %d", p_inst->inst_p_val[v]);
    printf("\n");
  }

  debug = 0;
#endif

}
#endif /* KERNEL_NBR_VAR */

void mcmc_main(Graph_t * p_graph, int num_insts, Instance_t * p_insts, 
               Task_t task, int end_time, double mem_limit, FILE * out_fp)
{
  int num_samples;
  double * p_weights;
  Stat_PR_t * p_pr_stats;
  Stat_MPE_t * p_mpe_stats;
  Stat_MAR_t * p_mar_stats;
  Stat_BEL_t * p_bel_stats;
  
  /* randomly initialize the starting state */
  mcmc_rand_init(p_graph, num_insts, p_insts);
  
  /* initialize memory for MCMC weights */
  p_weights = (double *) malloc(MAX_CARD * sizeof(*p_weights));
  
  /* initialize statistics */
  p_pr_stats = NULL;
  p_mpe_stats = NULL;
  p_mar_stats = NULL;
  p_bel_stats = NULL;
  
  if (Task_PR == task)
  {
    int i;
    
    p_pr_stats = (Stat_PR_t *)malloc(num_insts * sizeof(*p_pr_stats));
    
    for (i=0; i<num_insts; i++)
      stat_PR_init(p_pr_stats+i, p_graph);
  }
  else if (Task_MPE == task)
  {
    int i;
    
    p_mpe_stats = (Stat_MPE_t *)malloc(num_insts * sizeof(*p_mpe_stats));
    
    for (i=0; i<num_insts; i++)
      stat_MPE_init(p_mpe_stats+i, p_graph);
  }
  else if (Task_MAR == task)
  {
    int i;
    
    p_mar_stats = (Stat_MAR_t *)malloc(num_insts * sizeof(*p_mar_stats));
    
    for (i=0; i<num_insts; i++)
      stat_MAR_init(p_mar_stats+i, p_graph);
  }
  else if (Task_BEL == task)
  {
    int i;
    
    p_bel_stats = (Stat_BEL_t *)malloc(num_insts * sizeof(*p_bel_stats));
    
    for (i=0; i<num_insts; i++)
      stat_BEL_init(p_bel_stats+i, p_graph);
  }
  
  num_samples = 0;
  
  /* conservatively, finish one second before the due time */
  while (((int)time(NULL)) < (end_time - 1))
  {
    int i;
    /* for each instance */
    for (i=0; i<num_insts; i++)
    {
      Instance_t * p_inst;
      
      p_inst = p_insts + i;

      mcmc_kernel_1var(p_graph, p_inst, p_weights);

#ifdef KERNEL_NBR_VAR
      mcmc_kernel_nbr_var(p_graph, p_inst, p_weights);
#endif

      /* update statistics */
      if (Task_PR == task)
      {
        stat_PR_update(p_pr_stats+i, p_graph, p_inst->inst_p_val);
      }
      else if (Task_MPE == task)
      {
        stat_MPE_update(p_mpe_stats+i, p_graph, p_inst->inst_p_val);
      }
      else if (Task_MAR == task)
      {
        stat_MAR_update(p_mar_stats+i, p_graph, p_inst->inst_p_val);
      }
      else if (Task_BEL == task)
      {
        stat_BEL_update(p_bel_stats+i, p_graph, p_inst->inst_p_val);
      }
    }
    num_samples ++;
  }

  printf("%d samples\n", num_samples);

  fprintf(out_fp, "%d\n", num_insts);
  
  /* now output the solution */
  if (Task_PR == task)
  {
    int i;
    
    for (i=0; i<num_insts; i++)
      stat_PR_write(p_pr_stats+i, p_graph, out_fp);
    
    free(p_pr_stats);
  }
  else if (Task_MPE == task)
  {
    int i;
    
    for (i=0; i<num_insts; i++)
      stat_MPE_write(p_mpe_stats+i, p_graph, out_fp);
    
    free(p_mpe_stats);
  }
  else if (Task_MAR == task)
  {
    int i;
    
    for (i=0; i<num_insts; i++)
      stat_MAR_write(p_mar_stats+i, p_graph, out_fp);
    
    free(p_mar_stats);
  }
  else if (Task_BEL == task)
  {
    int i;
    
    for (i=0; i<num_insts; i++)
      stat_BEL_write(p_bel_stats+i, p_graph, out_fp);
    
    free(p_bel_stats);
  }

  free(p_weights);
}

